import jsonlines


def process_jsonl(input_file, output_train_file, output_val_file):
    with jsonlines.open(input_file) as reader, \
            jsonlines.open(output_train_file, mode='w') as train_writer, \
            jsonlines.open(output_val_file, mode='w') as val_writer:

        for obj in reader:
            if obj['split'] == 'train':
                train_writer.write(obj)
            elif obj['split'] == 'val':
                val_writer.write(obj)


process_jsonl('PubTabNet_2.0.0.jsonl',
              'PubTabNet_2.0.0_train.jsonl',
              'PubTabNet_2.0.0_val.jsonl')